# TCP middleman: forward a TCP connection between two peers
# and listen / modify traffic

import sys
import socket
import select
import binascii
import logging
import threading
import errno

# temporary imports
import random, time

from netutil import *
from rtmp import *                    # RTMP spec constants
from fo_client import RTMPStream, BufferedSocket, StreamingFailover
from fo_client import bind_socket_to_port_range
import fo_client

log = logging.getLogger('sim')
failoverlatency = open('failoverlatency.csv', 'w')

class SocketReader:
    '''
    Socket-reader that stores recently read data in self.recently_read.
    Warning: this has very bad performance unless self.recently_read is
    periodically cleared externally.
    '''

    def __init__(self, socket):
        self.recently_read = ''
        self.sock = socket

    def read(self, nbytes):
        data = ''
        while len(data) < nbytes:
            data += self.sock.recv(nbytes - len(data))
        self.recently_read += data
        return data

class FailoverSimulator:
    def __init__(self, cli_s, srv_addr, secs_til_failover):
        self.srv_addr = srv_addr
        self.cli_s = cli_s
        self.srv_s = None
        self.secs_til_failover = secs_til_failover
        self.srv_timeout = 2   # seconds to wait until server is assumed "dead"
        self.dump_traffic = False
        
    def simulate(self):
        try:
            self._simulate()
            log.info('simulation completed sucessfully')
        except socket.error as e:
            log.info('simulation ended from socket "error": %s', e)
        finally:
            if self.cli_s:
                self.cli_s.close()
            if self.srv_s:
                self.srv_s.close()
            log.info('exiting')

    def _simulate(self):
        "Listen on port srv_port for a connection and forward it to srv_addr"
        cli_addr = self.cli_s.getpeername()
        log.info('beginning simulation with client %s:%d', cli_addr[0], cli_addr[1])
        
        # Initiate connection with server
        self.srv_s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        
        # connect from port in range 50000-60000 -- iptables will
        # automatically drop this connection after 200 packets
        #port = bind_socket_to_port_range(self.srv_s, 50000, 60000)
        #log.debug("bound to port %d", port)
        self.srv_s.connect(self.srv_addr)
        log.debug("connected to server %s", self.srv_addr)

        # timeout-based sockets
        self.srv_s.settimeout(3)
        self.cli_s.settimeout(3)

        # forward data until a time limit is reached
        log.info("forwarding packets for %g seconds", self.secs_til_failover)
        self._forward_rtmp_chunks(secs=self.secs_til_failover)

        # forward data until connection is dropped (i.e. by iptables)
        '''log.info("forwarding packets indefinitely")
        try:
            self._forward()
        except socket.error as e:
            log.debug("socket error during forwarding : %s", e)
            if e.args[0] == errno.ECONNRESET:
                log.debug('server connection has been dropped')
            else:
                raise'''
        
        # close server connection
        log.debug('checking for remaining data from server')
        self.srv_s.settimeout(0)
        lastdata = ''
        try:
            lastdata = self.srv_s.recv(4096)
        except socket.error:
            pass
        log.debug('last data = %s\n%s', binascii.hexlify(lastdata), stringDump(lastdata))
        self.srv_s.close()
        self.srv_s = None
        log.debug("disconnected from server")

        log.debug('failoverlatency = %s', failoverlatency)
        failoverlatency.write('%s,%f,' % (threading.current_thread().name, time.time()))

        # inject get_state() RPC
        cmd = binascii.unhexlify("03" + "000000" + "000017" + "11" + "00000000" + "000200096765745f737461746500000000000000000005")
        self.cli_s.sendall(cmd)
        log.debug("sent get_state() to client")

        # wait for receive_state() response from client
        bs = BufferedSocket(self.cli_s)
        cli_RTMP = RTMPStream(bs)
        while 1:
            packet = cli_RTMP.read_packet()
            if packet.header.mtype == MTYPE_AMF3_COMMAND:
                log.debug('received packet, probably streaming_state')
                try:
                    streaming_state = packet.amf_payload[3]
                    log.debug("packet = %s", packet)
                    log.debug("amf_payload = %s", packet.amf_payload)
                    log.debug("streaming_state = %s", streaming_state)
                except (pyamf.DecodeError, IndexError) as e:
                    log.debug('packet not streaming state packet')
                    continue
                log.info('recreating state on server: %s', streaming_state)
                break
            else:
                log.debug('unsolicited packet from client:\n%s', packet)

        failoverlatency.write('%f,' % time.time())

        # recreate connection state on server side
        log.debug('beginning failover')
        f = StreamingFailover(self.srv_addr, streaming_state, port_range=(30000, 40000),
                              name='Failover #%s' % threading.current_thread().name[7:])
        f.failover()
        self.srv_s = f.sock

        failoverlatency.write('%f,' % time.time())
        
        # discard data before audio/video
        bs = BufferedSocket(self.srv_s)
        srv_RTMP = RTMPStream(bs)
        cli_RTMP = RTMPStream(BufferedSocket(self.cli_s))
        while 1:
            packet = srv_RTMP.read_packet()
            if packet.header.mtype in (MTYPE_AUDIO, MTYPE_VIDEO):
                #log.debug('forwarding first audio/video packet! (time = %f)',
                #          packet.header.time)
                #cli_RTMP.write_packet(packet)
                break
            else:
                log.debug('dropping connection setup packet: %s, %d', MTYPES.get(packet.header.mtype), packet.header.sid)

        # check the times on the next few packets...
        '''for i in range(30):
            packet = srv_RTMP.read_packet()
            if packet.header.mtype in (MTYPE_AUDIO, MTYPE_VIDEO):
                if packet.header.time > 1000 * streaming_state['time'] - 1:
                    log.debug('forwarding %s packet! (time = %f)',
                              MTYPES[packet.header.mtype], packet.header.time)
                    cli_RTMP.write_packet(packet)
                else:
                    log.debug('dropping old %s packet! (time = %f)',
                              MTYPES[packet.header.mtype], packet.header.time)
                    #break
            else:
                cli_RTMP.write_packet(packet)'''

        failoverlatency.write('%f\n' % time.time())

        log.debug('forwarding packets forever')
        
        # forward data between client and new server
        self.dump_traffic = True
        self.num_iters = 10
        self.srv_timeout = 1000000000
        self._forward(idle_time=100)
        log.info('exiting due to idle connection')

    def _forward(self, secs=None, idle_time=None):
        '''
        Forward available packets between client and server for secs seconds,
        or until connection is inactive for idle_time seconds.

        If secs is None, never stop forwarding unless due to inactivity.
        If idle_time is None, never stop forwarding due to inactivity.
        '''

        MSS = 4096
        now = time.time()
        idle_cutoff = idle_time is not None and (now + idle_time)
        end_time = secs is not None and (now + secs)
        last_srv_time = now

        while not end_time or time.time() < end_time:
            readable, writable, err = select.select((self.cli_s, self.srv_s),
                                                    (self.cli_s, self.srv_s),
                                                    (), 1)
            now = time.time()

            if idle_cutoff:
                if not readable:
                    if now >= idle_cutoff:
                        return
                else:
                    idle_cutoff = now + idle_time
        
            if self.cli_s in readable and self.srv_s in writable:
                data = self.cli_s.recv(MSS)
                self.srv_s.sendall(data)

            if self.srv_s in readable and self.cli_s in writable:
                self.data = self.srv_s.recv(MSS)
                self.cli_s.sendall(self.data)
                last_srv_time = now

            if now - last_srv_time > self.srv_timeout:
                log.debug('Server has been idle for %f seconds (now = %f, last_srv_time = %f, time.time() = %f)',
                          self.srv_timeout, now, last_srv_time, time.time())
                return

    def _forward_rtmp_chunks(self, secs):
        '''
        Forward available packets from server to client on RTMP
        chunk boundaries and from client to server irrespective
        of chunk boundaries for secs seconds.
        '''

        MSS = 4096
        end_time = secs + time.time()
        srv_reader = SocketReader(self.srv_s)
        to_cli_rtmp = RTMPStream(srv_reader)
        bytes_from_srv = 0

        log.debug("in _forward_rtmp_chunks")

        while time.time() < end_time:
            readable, writable, err = select.select((self.cli_s, self.srv_s),
                                                    (self.cli_s, self.srv_s),
                                                    (), 1)
            now = time.time()

            if self.cli_s in readable and self.srv_s in writable:
                data = self.cli_s.recv(MSS)
                self.srv_s.sendall(data)
                #log.debug("forward_rtmp_chunks: cli -> srv %d bytes", len(data))

            if self.srv_s in readable and self.cli_s in writable:
                # forward handshake bytes
                if bytes_from_srv < 3073:
                    data = self.srv_s.recv(3073 - bytes_from_srv)
                else:
                    srv_reader.recently_read = ''
                    chunk = to_cli_rtmp.read_chunk()
                    #log.debug('read chunk : %s', chunk)
                    data = srv_reader.recently_read
                bytes_from_srv += len(data)
                self.cli_s.sendall(data)
                #log.debug("forward_rtmp_chunks: srv -> cli %d bytes", len(data))


if __name__ == '__main__':
    #sys.argv[1:] = ['-s', 'sns9.cs.princeton.edu', '-l', '1980', '-v']
    sys.argv[1:] = ['-s', 'localhost', '-l', '1965', '-v']
    
    DEFAULT_TIME_TIL_FAILOVER = 6.3   # seconds until failover is carried out
    DEFAULT_LISTEN_PORT = 1965        # port on which to listen for RTMP clients
    DEFAULT_SERVER_ADDR = "localhost:1935"
    
    from optparse import OptionParser
    parser = OptionParser()
    parser.add_option("-t", "--time", type="int", dest="time", default=DEFAULT_TIME_TIL_FAILOVER,
                      help="time until failover (in seconds) [default %default]")
    parser.add_option("-r", "--random-time", action="store_true", dest="random_time",
                      help="randomize time until failover (3 to 12 seconds)")
    parser.add_option("-l", "--listen-port", dest="listen_port", default=DEFAULT_LISTEN_PORT,
                      type="int", help="incoming port to accept connections on [default %default]")
    parser.add_option("-s", "--server", dest="server_addr", default=DEFAULT_SERVER_ADDR,
                      help="address of server to forward connections to [default %default]")
    parser.add_option("-v", "--verbose", dest="verbose", action="store_true", default=False,
                      help="set logging output to debug-level [default %default]")    
    (options, args) = parser.parse_args()

    if ':' not in options.server_addr:
        options.server_addr += ':1935' # + str(RTMP_PORT)
    server_host, server_port = options.server_addr.split(':')
    server_addr = (server_host, int(server_port))

    logging_level = logging.INFO
    if options.verbose:
        logging_level = logging.DEBUG

    logging.basicConfig(level=logging.INFO, datefmt='%M:%S',
                        format='%(asctime)s.%(msecs)03d %(lineno)3d %(levelname)-5s %(threadName)s %(message)s')
    log.setLevel(logging_level)

    def failover_time():
        if options.random_time:
            return random.uniform(3, 12)
        else:
            return options.time

    failover_count = 0
    listen_s = cli_s = None
    try:
        listen_s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        listen_s.bind(('', options.listen_port))
        listen_s.listen(5)
        log.info('listening on port %d for RTMP connections', options.listen_port)
        while 1:
            failover_count += 1
            cli_s, cli_addr = listen_s.accept()
            log.debug("bound to client %s", cli_addr)
            sim = FailoverSimulator(cli_s, server_addr, failover_time())
            t = threading.Thread(None, sim.simulate, name="Thread %d" % failover_count)
            t.start()
        log.warning('FailoverSimulator is complete')
    finally:
        if listen_s:
            listen_s.close()
            log.info('closed listening socket (port %d)', options.listen_port)
        if cli_s: cli_s.close()
        failoverlatency.close()
